import torch
import torch.nn as nn
import torch.utils.data as Data
from torch.autograd import Variable
import numpy as np
import torchvision  # 数据库模块
import matplotlib.pyplot as plt

# 加cuda就是在训练数据、测试数据、神经网络部分加上cuda()
# print(torch.cuda.is_available())   #判断是否可使用cuda
# print(torch.cuda.device_count())  # 识别可以使用的gpu个数
# Hyper Parameters
EPOCH = 1  # 训练整批数据多少次, 为了节约时间, 我们只训练一次
BATCH_SIZE = 50  # 批大小
LR = 0.001  # 学习率
DOWNLOAD_MNIST = False  # 如果你已经下载好了mnist数据就写上 False

# Mnist 手写数字   训练集
train_data = torchvision.datasets.MNIST(
    root='./mnist/',  # 保存或者提取位置
    train=True,  # this is training data
    transform=torchvision.transforms.ToTensor(),  # 转换 PIL.Image or numpy.ndarray 成
    # torch.FloatTensor (C x H x W), 训练的时候 normalize 成 [0.0, 1.0] 区间
    download=DOWNLOAD_MNIST,  # 没下载就下载, 下载了就不用再下了
)

print(train_data.train_data.size())
print(train_data.train_labels.size())
plt.imshow(train_data.train_data[0].numpy(), cmap='gray')
plt.title('%i'%train_data.train_labels[0])
plt.show()

train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)

test_data = torchvision.datasets.MNIST(root='./mnist/', train=False)  # 测试集
test_x = torch.unsqueeze(test_data.data, dim=1).type(torch.FloatTensor)[
         :2000].cuda() / 255.  # shape from (2000, 28, 28) to (2000, 1, 28, 28), value in range(0,1)
test_y = test_data.targets[:2000].cuda()  # 测试集标签

# nn.Conv1d
# nn.MaxPool1d


# 定义卷积神经网络
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(  # input shape (1, 28, 28)
            nn.Conv2d(
                in_channels=1,  # input height
                out_channels=20,  # n_filters
                kernel_size=5,  # filter size
                stride=1,  # filter movement/step
                padding=2,  # 如果想要 con2d 出来的图片长宽没有变化, padding=(kernel_size-1)/2 当 stride=1
            ),  # output shape (16, 28, 28)
            nn.ReLU(),  # activation   --->(16, 28, 28)
            nn.MaxPool2d(kernel_size=2),  # 在 2x2 空间里向下采样, output shape   ---> (16, 14, 14)
        )
        self.conv2 = nn.Sequential(  # input shape  ---> (16, 14, 14)
            nn.Conv2d(20, 40, 5, 1, 2),  # output shape --->  (32, 14, 14)
            nn.ReLU(),  # activation
            nn.MaxPool2d(2),  # output shape   ---> (32, 7, 7)
        )
        self.out = nn.Linear(40 * 7 * 7, 10)  # fully connected layer, output 10 classes

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)  # ---> (batch,32, 7, 7)
        x = x.view(x.size(0), -1)  # 展平多维的卷积图成 (batch_size, 32 * 7 * 7)
        output = self.out(x)
        return output


cnn = CNN()
cnn.cuda()
# print(cnn)
optimizer = torch.optim.Adam(cnn.parameters(), lr=LR)  # optimize all cnn parameters
loss_func = nn.CrossEntropyLoss()  # the target label is not one-hotted

# training and testing
for epoch in range(EPOCH):
    for step, (x, y) in enumerate(train_loader):  # 分配 batch data, normalize x when iterate train_loader
        b_x = Variable(x).cuda()
        b_y = Variable(y).cuda()
        # vector_b = b_x.size()   [50,1,28,28]
        output = cnn(b_x)  # cnn output
        loss = loss_func(output, b_y)  # cross entropy loss
        optimizer.zero_grad()  # clear gradients for this training step
        loss.backward()  # backpropagation, compute gradients
        optimizer.step()  # apply gradients
        if step % 50 == 0:
            test_output = cnn(test_x)
            pred_y = torch.max(test_output, 1)[1].cuda().data.squeeze()
            right_sum = 0
            for a, b in zip(pred_y, test_y):
                if a == b:
                    right_sum += 1
                    accuracy = right_sum / test_y.size(0)
            print('Epoch:', epoch, '|step:', step, '|loss:', loss.data.cpu().numpy(), '|acc:', accuracy)

test_output = cnn(test_x[:10])
pred_y = torch.max(test_output, 1)[1].data.cpu().numpy().squeeze()
print(pred_y, 'prediction number')
print(test_y[:10].cpu().numpy(), 'real number')
